﻿//   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#include <math.h>
#include <iostream>
#include <string>
#include <vector>
#include "model_config.h"
#include "transforms.h"
#include "base_util/utils.h"

namespace ai {

std::map<std::string, int> interpolations = {{"LINEAR", cv::INTER_LINEAR},
                                             {"NEAREST", cv::INTER_NEAREST},
                                             {"AREA", cv::INTER_AREA},
                                             {"CUBIC", cv::INTER_CUBIC},
                                             {"LANCZOS4", cv::INTER_LANCZOS4}};

Transforms::Transforms(LogInfo* li):
    log_ifo(li) {

}

Transforms::~Transforms() {
  spdlog::get("logger")->info("Transforms::~Transforms");
}

bool Normalize::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  // static int img_idx = 0;
  // img_idx++;
  // if (img_idx % 50 == 0) {cv::imwrite(std::to_string(img_idx)+".jpg", img_blob.img);}
  // cv::imwrite(std::to_string(img_idx++)+".jpg", im);

  // 针对rk3588 不做任何处理 只需要赋值
  if (scale < 0) {
    data.img = im.clone();
    return true;
  }

  // onnx 处理方法
  float scale_value = 1.0;
  if (scale) { scale_value = 1./255.0; }
  if (mean_.empty()) {return false;}
  // std::cout << "scale_value:" << scale_value << std::endl;
  cv::Scalar mean = cv::Scalar(mean_[0], mean_[1], mean_[2]);
  if (cfg->data_format == "HWC") {
    if (std_.empty()) {
	    data.img = cv::dnn::blobFromImage(im, scale_value, im.size(), mean, false, false);
    } else {
      // rockchip video cls 处理方法
      im.convertTo(im, CV_32F, scale_value);
      cv::subtract(im, cv::Scalar(mean_[0], mean_[1], mean_[2]), im);
      cv::divide(im, cv::Scalar(std_[0], std_[1], std_[2]), im);
      data.img = im;
    }
  }
  // onnx 处理方法
  else if (cfg->data_format == "CHW") {
    if (std_.empty()) {return false;}
    // data.img = cv::dnn::blobFromImage(im, scale_value, im.size(), mean, true);
    int h = im.rows;
    int w = im.cols;
    int c = im.channels();
    im.convertTo(im, CV_32F, scale_value);
    cv::subtract(im, cv::Scalar(mean_[0], mean_[1], mean_[2]), im);
    cv::divide(im, cv::Scalar(std_[0], std_[1], std_[2]), im);
    // if (log_ifo->log_core_img) spdlog::get("logger")->info("CHW: [{},{},{}]", c, h, w);
    data.im_vec_data.resize(c * h * w);
    float* ptr = data.im_vec_data.data();
    for (int i = 0; i < c; ++i) {
      cv::extractChannel(im, cv::Mat(h, w, CV_32FC1, ptr + i * h * w), i);
    }
  }

  // std::vector<cv::Mat> split_im;
  // cv::split(*im, split_im);
  // #pragma omp parallel for num_threads(im.channels())
  // for (int c = 0; c < im.channels(); c++) {
  //   float range_val = max_val_[c] - min_val_[c];
  //   // cv::subtract(split_im[c], cv::Scalar(min_val_[c]), split_im[c]);
  //   // cv::divide(split_im[c], cv::Scalar(range_val), split_im[c]);
  //   // cv::subtract(split_im[c], cv::Scalar(mean_[c]), split_im[c]);
  //   // cv::divide(split_im[c], cv::Scalar(std_[c]), split_im[c]);
  // }
  // cv::merge(split_im, *im);
  // spdlog::get("logger")->info("Normalize::run {}", im.channels());

  // float scale_value = 1.0;
  // if (scale) {
  //   scale_value = 1./255.0;
  // }
  // im.convertTo(im, CV_32FC(im.channels()), scale_value);
  // cv::subtract(im, cv::Scalar(mean_[0], mean_[1], mean_[2]), im);
  // cv::divide(im, cv::Scalar(std_[0], std_[1], std_[2]), im);

  // int h = im.rows;
  // int w = im.cols;
  // int c = im.channels();
  // // if (log_ifo->log_core_img) spdlog::get("logger")->info("CHW: [{},{},{}]", c, h, w);
  // data.im_vec_data.resize(c * h * w);
  // float* ptr = data.im_vec_data.data();
  // for (int i = 0; i < c; ++i) {
  //   cv::extractChannel(im, cv::Mat(h, w, CV_32FC1, ptr + i * h * w), i);
  // }

  // for (int i = 0; i < c; ++i) {
  //   memcpy(ptr + i * h * w, split_im[i].data, h * w * sizeof(float));
  // }
  // spdlog::get("logger")->info("Normalize::run ok.");

  return true;
}

float ResizeByShort::GenerateScale(const cv::Mat& im) {
  int origin_w = im.cols;
  int origin_h = im.rows;
  int im_size_max = std::max(origin_w, origin_h);
  int im_size_min = std::min(origin_w, origin_h);
  float scale =
      static_cast<float>(short_size_) / static_cast<float>(im_size_min);
  if (max_size_ > 0) {
    if (round(scale * im_size_max) > max_size_) {
      scale = static_cast<float>(max_size_) / static_cast<float>(im_size_max);
    }
  }
  return scale;
}

bool ResizeByShort::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  data.im_shape_record.push_back({im.rows, im.cols});
  data.deal_order.push_back("ResizeByShort");

  float scale = GenerateScale(im);
  int width = static_cast<int>(round(scale * im.cols));
  int height = static_cast<int>(round(scale * im.rows));
  cv::resize(im, im, cv::Size(width, height), 0, 0, cv::INTER_LINEAR);

  data.new_im_shape = {im.rows, im.cols};
  data.scale = scale;
  return true;
}

bool CenterCrop::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  int height = static_cast<int>(im.rows);
  int width = static_cast<int>(im.cols);
  if (height < height_ || width < width_) {
    std::cerr << "[CenterCrop] Image size less than crop size" << std::endl;
    return false;
  }
  int offset_x = static_cast<int>((width - width_) / 2);
  int offset_y = static_cast<int>((height - height_) / 2);
  cv::Rect crop_roi(offset_x, offset_y, width_, height_);
  im = im(crop_roi);
  data.new_im_shape = {im.rows, im.cols};
  return true;
}

void Padding::GeneralPadding(cv::Mat& im,
                             const std::vector<float> &padding_val,
                             int padding_w, int padding_h) {
  cv::Scalar value;
  if (im.channels() == 1) {
    value = cv::Scalar(padding_val[0]);
  } else if (im.channels() == 2) {
    value = cv::Scalar(padding_val[0], padding_val[1]);
  } else if (im.channels() == 3) {
    value = cv::Scalar(padding_val[0], padding_val[1], padding_val[2]);
  } else if (im.channels() == 4) {
    value = cv::Scalar(padding_val[0], padding_val[1], padding_val[2],
                                  padding_val[3]);
  }

  // int top = padding_h / 2;
  // int bottom = padding_h / 2 + padding_h % 2 == 0 ? 0 : 1;
  // int left = padding_w / 2;
  // int right = padding_w / 2 + padding_w % 2 == 0 ? 0 : 1;
  // cv::copyMakeBorder(
  // im,
  // im,
  // top,
  // bottom,
  // left,
  // right,
  // cv::BORDER_CONSTANT,
  // value);
  cv::copyMakeBorder(
  im,
  im,
  0,
  padding_h,
  0,
  padding_w,
  cv::BORDER_CONSTANT,
  value);
}

void Padding::MultichannelPadding(cv::Mat& im,
                                  const std::vector<float> &padding_val,
                                  int padding_w, int padding_h) {
  std::vector<cv::Mat> padded_im_per_channel(im.channels());
  #pragma omp parallel for num_threads(im.channels())
  for (size_t i = 0; i < im.channels(); i++) {
    const cv::Mat per_channel = cv::Mat(im.rows + padding_h,
                                        im.cols + padding_w,
                                        CV_32FC1,
                                        cv::Scalar(padding_val[i]));
    padded_im_per_channel[i] = per_channel;
  }
  cv::Mat padded_im;
  cv::merge(padded_im_per_channel, padded_im);
  cv::Rect im_roi = cv::Rect(0, 0, im.cols, im.rows);
  im.copyTo(padded_im(im_roi));
  im = padded_im;
}

bool Padding::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  data.im_shape_record.push_back({im.rows, im.cols});
  data.deal_order.push_back("Padding");

  int padding_w = 0;
  int padding_h = 0;
  if (width_ > 1 & height_ > 1) {
    padding_w = width_ - im.cols;
    padding_h = height_ - im.rows;
  } else if (coarsest_stride_ >= 1) {
    int h = im.rows;
    int w = im.cols;
    padding_h =
        ceil(h * 1.0 / coarsest_stride_) * coarsest_stride_ - im.rows;
    padding_w =
        ceil(w * 1.0 / coarsest_stride_) * coarsest_stride_ - im.cols;
  }

  if (padding_h < 0 || padding_w < 0) {
    std::cerr << "[Padding] Computed padding_h=" << padding_h
              << ", padding_w=" << padding_w
              << ", but they should be greater than 0." << std::endl;
    return false;
  }
  if (im.channels() < 5) {
    Padding::GeneralPadding(im, im_value_, padding_w, padding_h);
  } else {
    Padding::MultichannelPadding(im, im_value_, padding_w, padding_h);
  }
  // spdlog::get("logger")->info("Padding::run {}, {}", im.cols, im.rows);

  data.new_im_shape = {im.rows, im.cols};
  return true;
}

bool ResizeByLong::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  data.im_shape_record.push_back({im.rows, im.cols});
  data.deal_order.push_back("ResizeByLong");
  int origin_w = im.cols;
  int origin_h = im.rows;
  // spdlog::get("logger")->info("ResizeByLong::run {}, {}", origin_w, origin_h);

  float scale = 1.0;
  int im_size_max = std::max(origin_w, origin_h);
  if (fix_size > 0) { 
    scale = fix_size * 1.0 / im_size_max;
  }
  if (max_size > 0) {
    if (im_size_max < max_size) {
      scale = 1.0;
    } else {
      scale = max_size * 1.0 / im_size_max;
    }
  }

  cv::resize(im, im, cv::Size(), scale, scale, cv::INTER_NEAREST);

  data.new_im_shape = {im.rows, im.cols};
  data.scale = scale;
  return true;
}

bool Resize::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  if (interpolations.count(interp_) <= 0) {
    std::cerr << "[Resize] Invalid interpolation method: '" << interp_ << "'"
              << std::endl;
    return false;
  }
  data.im_shape_record.push_back({im.rows, im.cols});
  data.deal_order.push_back("Resize");

  float scale = 1.0;
  if (target_size.size() > 0) {
    scale = std::min(target_size[1] * 1.0 / im.rows, target_size[0] * 1.0 / im.cols);
    int dst_width = int(scale * im.cols);
    int dst_height = int(scale * im.rows);
    cv::resize(im, im, cv::Size(dst_width, dst_height),0.,0., cv::INTER_NEAREST);

    int bottom = target_size[1] - dst_height;
    int right = target_size[0] - dst_width;
    if (bottom < 0 || right < 0) { return false; }
    cv::copyMakeBorder(im, im, 0, bottom, 0, right, cv::BORDER_CONSTANT);
  }
  else if (fix_size.size() > 0) {
    cv::resize(im, im, cv::Size(fix_size[0], fix_size[1]), 0., 0., cv::INTER_NEAREST);
  }


  data.new_im_shape = {im.rows, im.cols};
  data.scale = scale;
  return true;
}


bool ResizeText::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  float wh_ratio = float(im.cols) / float(im.rows);;
  int imgW = int(ceilf(std_height * wh_ratio));

  data.im_shape_record.push_back({im.rows, im.cols});
  data.deal_order.push_back("ResizeText");

  cv::resize(im, im, cv::Size(imgW, std_height), 0.f, 0.f, interpolations[interp_]);
  data.new_im_shape = {im.rows, im.cols};
  return true;
}

bool Clip::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  std::vector<cv::Mat> split_im;
  cv::split(im, split_im);
  for (int c = 0; c < im.channels(); c++) {
    cv::threshold(split_im[c], split_im[c], max_val_[c], max_val_[c],
                  cv::THRESH_TRUNC);
    cv::subtract(cv::Scalar(0), split_im[c], split_im[c]);
    cv::threshold(split_im[c], split_im[c], min_val_[c], min_val_[c],
                  cv::THRESH_TRUNC);
    cv::divide(split_im[c], cv::Scalar(-1), split_im[c]);
  }
  cv::merge(split_im, im);
  return true;
}

int Transforms::init(const YAML::Node& transforms_node) {
  transforms.clear();
  std::string name;
  for (const auto& item : transforms_node) {
    name = item.begin()->first.as<std::string>();
    if (name == "ArrangeClassifier") { continue; }
    if (name == "ArrangeSegmenter") { continue; }
    if (name == "ArrangeFasterRCNN") { continue; }
    if (name == "ArrangeMaskRCNN") { continue; }
    if (name == "ArrangeYOLOv3") { continue; }

    Transform* transform = CreateTransform(name);
    transform->init(item.begin()->second);
    transforms.push_back(transform);
  }
  //最后一个必须是 Normalize
  if (name != "Normalize") {return -1;}
  return 0;
}

Transform* Transforms::CreateTransform(
    const std::string& transform_name) {
  if (transform_name == "Normalize") {
    return new Normalize();
  } else if (transform_name == "ResizeByShort") {
    return new ResizeByShort();
  } else if (transform_name == "CenterCrop") {
    return new CenterCrop();
  } else if (transform_name == "Resize") {
    return new Resize();
  } else if (transform_name == "ResizeText") {
    return new ResizeText();
  } else if (transform_name == "Padding") {
    return new Padding();
  } else if (transform_name == "ResizeByLong") {
    return new ResizeByLong();
  } else if (transform_name == "Clip") {
    return new Clip();
  } else {
    std::cerr << "There's unexpected transform(name='" << transform_name
              << "')." << std::endl;
    exit(-1);
  }
}

bool Transforms::run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) {
  // do all preprocess ops by order
  // printf("Transforms::run start\n");

  // 20250702 撤销拷贝操作 降低内存占用
  // cv::Mat img = im.clone();
  if (log_ifo->log_level_5) { spdlog::get("logger")->info("img.channels(): {}", im.channels()); }

  if (im.channels() == 1) { cv::cvtColor(im, im, cv::COLOR_GRAY2BGR); }
  if (cfg->img_type == std::string("RGB")) { cv::cvtColor(im, im, cv::COLOR_BGR2RGB); }

  data.ori_im_shape = {im.rows, im.cols};
  for (int i = 0; i < transforms.size(); ++i) {
    double t1 = 0;
    if (log_ifo->log_level_4) { t1 = cv::getTickCount(); }
    // 最后一个必须是 Normalize
    if (!transforms[i]->run(im, data, cfg)) {
      spdlog::get("logger")->info("Apply transforms to image failed!");
      return false;
    }

    if (log_ifo->log_level_4) {
      t1 = ((double)cv::getTickCount() - t1) / cv::getTickFrequency();
      spdlog::get("logger")->info("transforms {} Run Time: {}", i, t1);
    }
  }

  // data.new_im_shape = {img.rows, img.cols};

  // // data format NHWC to NCHW
  // // img data save to ImageBlob
  // int h = img.rows;
  // int w = img.cols;
  // int c = img.channels();
  // if (log_ifo->log_core_img) spdlog::get("logger")->info("CHW: [{},{},{}]", c, h, w);
  // data->im_data.resize(c * h * w);
  // float* ptr = data->im_data.data();
  // for (int i = 0; i < c; ++i) {
  //   cv::extractChannel(img, cv::Mat(h, w, CV_32FC1, ptr + i * h * w), i);
  // }
  
  if (log_ifo->log_level_4 && !data.im_vec_data.empty()) {
    auto res = std::minmax_element(data.im_vec_data.begin(), data.im_vec_data.end());
    spdlog::get("logger")->info("DEBUG min: {} max: {}", *(res.first), *(res.second));
  }
  return true;
}

}  // namespace PaddleX
